(* simplifies the murphi ast *)

open Ast
open Mods

type error =
    Unsupported_multdiv
  | Variables_addsub
  | Array_index
  | Undefined_type of string
  | Undefined_var of string
  | Undefined_fun of string
  | Undefined_field of string
  | Not_an_array
  | Not_a_record
  | Designator_usage
  
exception Error of error 

type desig_use = DUarray_index | DUexpr
    
let types = Hashtbl.create 32 (* (id, type_expr) *)
let vars = Hashtbl.create 32 (* (id,  type_expr) *)
let funs = Hashtbl.create 32 (* (id,  (type_expr, type_expr StrMap.t) ) *)

let desigs = Hashtbl.create 32 (* (designator.node, desig_use) *)
  
let mk_node e t = {data = t; node = e}
let mk_any e = {data = Tany; node = e}
  
let loc_dummy e =
  mk_node e (Lexing.dummy_pos, Lexing.dummy_pos)
      
let rec prune_quantif q =
  mk_node (
           match q.node with
           | Qvar (id, te) -> Qvar (id, prune_type_expr te)
           | Qfor (id, e1, e2, e3) ->
               Qfor (id, prune_expr e1, prune_expr e2, prune_expr e3)
  ) Tany

and prune_quantifs ql =
  List.map prune_quantif ql

and designator_usage du d =
  try
    let usage = Hashtbl.find desigs d in
    if du <> usage then raise (Error Designator_usage)
  with Not_found -> Hashtbl.add desigs d du

    
and prune_desig designator =
  match designator.node with
  | DESid id -> mk_node (DESid id)
      begin try Hashtbl.find vars id.node
       with Not_found -> raise (Error (Undefined_var id.node)) end
  | DESarray (d, e) ->
      let pd = prune_desig d in
      mk_node (DESarray (
        pd,
        begin match e.node with
        | Edesignator ed ->
            let e = begin match ed.node with 
            | DESdot (_, _) -> prune_expr e
            | (DESid id) -> prune_expr e
            | _ -> raise (Error Array_index)
            end in
            designator_usage DUarray_index ed.node;
            e
        | _ -> raise (Error Array_index)
        end
      ))
        begin match pd.data with | Tarray (t, _) -> t
        | _ -> raise (Error Not_an_array) end
  | DESdot (d, id) ->
      let pd = prune_desig d in
      let dt = begin match pd.data with
      | Tident ident ->
          let t = begin try Hashtbl.find types ident.node
            with Not_found -> raise (Error (Undefined_type ident.node)) end
          in
          begin match t with Trecord m -> begin try StrMap.find id.node m with
            Not_found -> raise (Error (Undefined_field id.node)) end
          | _ -> raise (Error Not_a_record) end
      | Trecord m -> begin try StrMap.find id.node m with
            Not_found -> raise (Error (Undefined_field id.node)) end
      | _ -> raise (Error Not_a_record)
      end in
      mk_node (DESdot (prune_desig d, id)) dt
      
(* TODO : mettre le bon type une fois que David aura fini de modifier
l'ast et le parser*) 

and prune_stmt stmt =
  mk_node (match stmt.node with
  | Sassign (d, e) -> Sassign (prune_desig d, prune_expr e)
  | Sif (e, sl1, sl2) -> Sif (prune_expr e, prune_stmts sl1, prune_stmts sl2)
  | Sswitch (e, esl, sl) ->
      Sswitch (prune_expr e,
               List.map (fun (el, sl) ->
                 (prune_exprs el, prune_stmts sl)) esl, prune_stmts sl)        
  | Sfor (ql, sl) -> Sfor (prune_quantifs ql, prune_stmts sl)
  | Swhile (e, sl) -> Swhile (prune_expr e, prune_stmts sl)
  | Salias (al, sl) ->
      Salias (List.map (fun (id, e) -> (id, prune_expr e)) al, prune_stmts sl)
  | Sproccall (id, el) -> Sproccall (id, prune_exprs el)
  | Sundef d -> Sundef (prune_desig d)
  | Serror s -> Serror (s)
  | Sclear d -> Sclear (prune_desig d)
  | Sassert (e, os) -> Sassert (prune_expr e, os)
  | Sput (Putexpr pe) -> Sput (Putexpr (prune_expr pe))
  | Sput (Putstring _) as s -> s
           (** Correction sur les deux lignes au dessus **)
  | Sreturn (None) -> Sreturn (None)
  | Sreturn (Some e) -> Sreturn (Some (prune_expr e))
  | Smultisetadd (d1, d2) -> Smultisetadd (prune_desig d1, prune_desig d2)
  | Smultisetremovepred (id, d, e) ->
      Smultisetremovepred (id, prune_desig d, prune_expr e)
  | Smultisetremove (e, d) -> Smultisetremove (prune_expr e, prune_desig d)
  ) Tany

and prune_stmts sl =
  List.map prune_stmt sl

and prune_expr e =
  prune_expr_ex false e

and prune_expr_ex is_cst e =
  match e.node with
  | Eforall (ql, e) -> mk_node (Eforall (prune_quantifs ql, prune_expr e)) Tany
  | Eexists (ql, e) -> mk_node (Eexists (prune_quantifs ql, prune_expr e)) Tany
  | Ebinop (op, e1, e2) ->
      mk_node (begin match op with
      | Bplus | Bminus ->
          begin match (e1.node, e2.node) with
            | (Edesignator d1, Edesignator d2) ->
                designator_usage DUexpr d1.node;
                designator_usage DUexpr d2.node
            | (_, Edesignator d) | (Edesignator d, _) ->
                designator_usage DUexpr d.node
            | _ -> ()
          end;
          begin match (e1.node, e2.node) with
            | (_, Econst _) | (Econst _, _) -> ()
            | (_, _) -> if (not is_cst) then raise (Error Variables_addsub)
              else ()
          end
      | _ -> raise (Error Unsupported_multdiv)
      end;
      Ebinop (op, prune_expr e1, prune_expr e2)) Tint  (* OK, on sait faire sous Cubicle *)
  | Enot e -> mk_node (Enot (prune_expr e)) (Tident (loc_dummy "boolean"))
  | Econdition (e1, e2, e3) ->
      mk_node (Econdition (prune_expr e1, prune_expr e2, prune_expr e3))
        (Tident (loc_dummy "boolean"))
  | Edesignator d -> let pd = prune_desig d in mk_node (Edesignator pd) pd.data
                    (** Correction sur la ligne du dessus **)
  | Econst i32 -> mk_node (Econst i32) Tint
  | Efuncall (id, el) ->
      mk_node (Efuncall (id, prune_exprs el))
        (let t, _ = try Hashtbl.find funs id
          with Not_found -> raise (Error (Undefined_fun id.node)) in t)
  | Eerror s -> mk_node (Eerror s) Tany
  | Emultiset_count (id, d, e) ->
      mk_node (Emultiset_count (id, prune_desig d, prune_expr e)) Tint
  | Eismem (d, te) ->
      mk_node (Eismem (prune_desig d, prune_type_expr te))
        (Tident (loc_dummy "boolean"))
  | Eisundef d -> mk_node (Eisundef (prune_desig d))
      (Tident (loc_dummy "boolean"))
                               
and prune_exprs el =
  List.map prune_expr el

and prune_type_expr = function 
  | Tident (id) -> (* Tident (id) *)
      let t = try Hashtbl.find types id.node with
        Not_found -> raise (Error (Undefined_type id.node))
      in t
  | Tsubrange (i1, i2) -> Tsubrange (i1, i2) 
  | Tenum (idl) -> Tenum (idl)
  | Trecord (vdl) as s -> s (** A priori c'est déjà OK, le parser fait le job **)
(*      Trecord (List.map (fun (id, te) -> (id, prune_type_expr te)) vdl)*)
  | Tarray (te1, te2) -> Tarray (prune_type_expr te1, prune_type_expr te2)
  | Tscalarset i -> Tscalarset i
  | Tunion te -> Tunion (prune_type_exprs te)
  | Tmultiset (e, te) -> Tmultiset (e, prune_type_expr te)
  | Tint -> Tint
  | Tany -> Tany

and prune_type_exprs tel =
  List.map prune_type_expr tel

let prune_proc_decl var = (*match var with
  | _ as pd -> pd*) failwith "TODO"

let prune_proc_decls pdl =
  List.map prune_proc_decl pdl

let rec prune_rule rule =
  mk_node (match rule.node with
  | Rsimple (os, e, dl, sl) -> Rsimple (os, prune_expr e, prune_decls dl, prune_stmts sl)
  | Rstart (os, dl, sl) -> Rstart (os, prune_decls dl, prune_stmts sl)
  | Rinvariant (os, e) -> Rinvariant (os, prune_expr e)
  | Rruleset (ql, rl) -> Rruleset (prune_quantifs ql, prune_rules rl)
  | Ralias (al, rl) -> Ralias (List.map (fun (id, e) -> (id, prune_expr e)) al, prune_rules rl)
  | Rchoose (id, d, rl) -> Rchoose (id, prune_desig d, prune_rules rl))
  Tany

and prune_rules rl =
  List.map prune_rule rl

and prune_decl (var : Ast.loc Ast.decl) : Ast.type_expr Ast.decl = match var with
(*  | Dconst (id, e) ->
      let pe = prune_expr_ex true e in
      Hashtbl.add vars id.node pe.data;
      Dconst (id, pe)*)
  | Dtype (id, te) ->
      let pte = prune_type_expr te in
      Hashtbl.add types id.node pte;
      Dtype (id, pte)
  | Dvar (id, te) ->
      let pte = prune_type_expr te in
      let pte = begin match pte with
        (** Quand var d'un type nommé, rattacher directement **)
        (** le type, sauf si c'est un type Tenum **)
      | Tident id -> let tpte = begin try Hashtbl.find types id.node
            with Not_found -> raise (Error (Undefined_type id.node)) end
          in begin match tpte with
          | Tenum _ -> pte
          | _ -> tpte
          end
        (** Quand var est de type Tenum, lui donner un nom **)
        (** est lui donner à la var le type Tident **)
      | Tenum _ ->
           let tname = Printf.sprintf "%s_type" id.node in
           Hashtbl.add types tname pte;
           Tident (loc_dummy tname)
      | _ -> pte
      end in
      Hashtbl.add vars id.node pte;
      Dvar (id, pte)

and prune_decls (dl : Ast.loc Ast.decl list) : Ast.type_expr Ast.decl list =
  List.map prune_decl dl

let prune_prog (p : Ast.loc Ast.prog) : Ast.type_expr Ast.prog =
  Hashtbl.add types "boolean"
    (Tenum [loc_dummy "True"; loc_dummy "False"]);
  let (dl, pfdl, rl) = p in
  let pdl = prune_decls dl in
  let ppfdl = prune_proc_decls pfdl in
  (*let prl = prune_rules rl in*)
  let prl = [] in
  (pdl, ppfdl, prl)

